import torch

from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split

class DatasetSpliter:
    def __init__(self, dataset, batch_size):
        self.dataset = dataset
        self.batch_size = batch_size

    def split(self):
        train_indices, test_indices = train_test_split(range(len(self.dataset)), test_size=0.1, random_state=1219)  # 4891
        train_dataset = torch.utils.data.Subset(self.dataset, train_indices)
        test_dataset = torch.utils.data.Subset(self.dataset, test_indices)
        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
        test_loader = DataLoader(test_dataset, batch_size=self.batch_size, shuffle=False)
        log_text = f"Dataset prepared, dataset length: {len(self.dataset.labels)}"
        print(log_text)

        return train_loader, test_loader
